# ！/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/1/8 22:07
# @Author  : Zhang He
# @FileName: pdf_merge.py
# @Software: PyCharm

import os
import PyPDF2


def merge_pdf(input_path):
    output_path = r"output/merged_pdf.pdf"
    files = [file for file in os.listdir(input_path) if file.endswith(".pdf")]
    print(files)
    merger = PyPDF2.PdfMerger()
    for filename in files:
        merger.append(PyPDF2.PdfReader(os.path.join(input_path, filename)))
    merger.write(output_path)


def delete_page(input_path, del_page_list):
    # 获取PDF文件名称
    files = [file for file in os.listdir(input_path) if file.endswith(".pdf")]
    original_pdf = os.path.join(input_path, files[0])
    print(original_pdf)

    # 读取原始的PDF文件
    pdf_reader = PyPDF2.PdfReader(original_pdf)

    # 获取原始文件中的总页数
    total_pages = len(pdf_reader.pages)

    # 指定要删除的页面索引或范围，从0开始计数
    delete_pages = del_page_list

    # 创建一个新的PdfFileWriter对象
    pdf_writer = PyPDF2.PdfWriter()

    # 循环遍历每一页
    for i in range(total_pages):
        # 如果当前页索引不在要删除的页面列表中，则保留该页
        if i not in delete_pages:
            # 获取当前页对象
            page = pdf_reader.pages[i]
            # 将当前页对象添加到PdfFileWriter对象中
            pdf_writer.add_page(page)

    # 创建一个新的PDF文件名，格式为"original_deleted.pdf"
    new_file = r"output/deleted_pdf.pdf"

    # 打开一个新的PDF文件，以二进制写入模式
    with open(new_file, "wb") as f:
        # 将PdfFileWriter对象中的内容写入到新的PDF文件中
        pdf_writer.write(f)

    # 打印出新创建的PDF文件名
    print(new_file)


def compress(text, quality):
    # 使用Python内置的gzip库进行压缩
    import gzip
    # 将文本转换为字节对象并压缩
    compressed_bytes = gzip.compress(text.encode('utf-8'))
    # 将压缩后的字节对象转换回字符串对象
    compressed_str = compressed_bytes.decode('utf-8', 'ignore')
    # 返回压缩后的字符串对象
    return compressed_str


def compress_pdf(input_file, output_file, quality):
    # 打开输入文件
    pdf_file = open(input_file, 'rb')
    # 创建PDF阅读器对象
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    # 创建PDF写入器对象
    pdf_writer = PyPDF2.PdfWriter()
    # 遍历每一页并重新编码
    for page_num in range(len(pdf_reader.pages)):
        page = pdf_reader.pages[page_num]
        # 重新编码页面
        # page_str = page.extract_text()
        # compressed_str = compress(page_str, quality)
        # compressed_bytes = compressed_str.encode('utf-8')
        # compressed_page = PyPDF2.PageObject.createTextString(compressed_bytes)

        # 将压缩后的页面添加到PDF写入器对象中
        # pdf_writer.add_page(compressed_page)

        page.compress_content_streams()
        pdf_writer.add_page(page)


    # 将压缩后的PDF写入输出文件
    with open(output_file, 'wb') as output:
        pdf_writer.write(output)
    # 关闭输入文件和输出文件
    pdf_file.close()
    output.close()


if __name__ == "__main__":
    # input_path = r"./input/"
    # output_path = r"output/merged_pdf.pdf"

    # 合并多个pdf文件
    # merge_pdf(input_path)  # 合并pdf

    # 删除指定pdf页面
    # del_page_list = [3]
    # del_page_list = [i for i in range(4, 9, 1)]
    # del_page_list = [i for i in range(18, 49, 1)] + [i for i in range(60, 78, 1)]
    # print(del_page_list)
    # delete_page(input_path, del_page_list)

    # 压缩pdf文件【以下程序还没有明显压缩效果】
    compress_pdf(input_file=r"input/merged_pdf.pdf", output_file=r"output/compressed_pdf.pdf", quality=0.8)


